package liangtcp // 提供消除了粘包之后消息

import (
	//"encoding/json"
	"bytes"
	"encoding/binary"
	"log"
	"net"

	"gitee.com/viplianghl/golib/liangaes"
)

type OnMsgFunc func(net.Conn, []byte) error
type OnConnFunc func(net.Conn) error
type OnDisconnFunc func(net.Conn) error

type LiangTCP interface {
	SendMsg(conn net.Conn, buf []byte) error
}

// -------------------------------------------------------------LiangTCPServer
type LiangTCPServer struct {
	OnMsg     OnMsgFunc
	OnConn    OnConnFunc
	OnDisconn OnDisconnFunc
	//port  int
	AESKey []byte
}

func (t *LiangTCPServer) SendMsg(conn net.Conn, buf []byte) error {

	// see AES encrpt
	if t.AESKey != nil {
		buf1, err2 := liangaes.AesCbcNoSaltEncrypt(buf, t.AESKey, t.AESKey)
		if err2 != nil {
			log.Println("LiangTCP SendMsg encrypt err: ", err2)
			return err2
		}
		buf = buf1
	}

	// write len
	temp := make([]byte, 4)
	binary.BigEndian.PutUint32(temp, uint32(len(buf)))
	n, err := conn.Write(temp)
	if err != nil || n != 4 {
		log.Println("LiangTCP SendMsg err: ", err)
		return err
	}

	// write data
	n1, err1 := conn.Write(buf)
	if err1 != nil || n1 != len(buf) {
		log.Println("LiangTCP SendMsg err: ", err)
		return err
	}
	return nil
}

func (t *LiangTCPServer) Listen(addr string) *LiangTCPServer {

	l, err := net.Listen("tcp", addr)
	if err != nil {
		log.Println("LiangTCP listen err: ", err)
		return t
	}
	defer l.Close()
	log.Println("LiangTCP listen at: ", addr)

	for {
		conn, err := l.Accept()
		if err != nil {
			log.Println("LiangTCP accept err: ", err)
			break
		}
		go t.handleConn(conn)
	}
	return t
}

func (t *LiangTCPServer) handleConn(conn net.Conn) {
	defer conn.Close()
	log.Println("LiangTCP A new client connected: ", conn.RemoteAddr().String())
	t.OnConn(conn)

	// 粘包处理
	var msgBuffer bytes.Buffer

	for {
		// 单次接收
		buf := make([]byte, 2000)
		n, err := conn.Read(buf)
		if err != nil {
			log.Printf("LiangTCP read error(from %s): %s\r\n", conn.RemoteAddr().String(), err)
			break
		}
		// log.Printf("LiangTCP read %d bytes(from %s\r\n", n, conn.RemoteAddr().String())

		msgBuffer.Write(buf[:n])
		err1 := t.handleMsg(conn, &msgBuffer)
		if err1 != nil {
			log.Printf("LiangTCP protocol error(from %s): %s\r\n", conn.RemoteAddr().String(), err)
			break
		}
	}
	t.OnDisconn(conn)
}

func (t *LiangTCPServer) handleMsg(conn net.Conn, msg *bytes.Buffer) error {
	for {
		if msg.Len() < 4 {
			// not enough for head length
			return nil
		}

		length := binary.BigEndian.Uint32(msg.Bytes())
		if int(length) > (msg.Len() - 4) {
			// not enough data
			return nil
		}

		// prepare data
		msg.Next(4)
		buf := msg.Next(int(length))

		// see AES decrypt
		if t.AESKey != nil {
			buf1, err2 := liangaes.AesCbcNoSaltDecrypt(buf, t.AESKey, t.AESKey)
			if err2 != nil {
				log.Println("LiangTCP RecvMsg decrypt err: ", err2)
				return err2
			}
			buf = buf1
		}

		// data is ready, call OnMsg
		err := t.OnMsg(conn, buf)
		if err != nil {
			return err
		}
	}
}

// -------------------------------------------------------------LiangTCPClient
type LiangTCPClient struct {
	OnMsg     OnMsgFunc
	OnDisconn OnDisconnFunc
	//port  int
	AESKey []byte
	conn   net.Conn
}

func (t *LiangTCPClient) SendMsg(buf []byte) error {

	// see AES encrypt
	if t.AESKey != nil {
		buf1, err2 := liangaes.AesCbcNoSaltEncrypt(buf, t.AESKey, t.AESKey)
		if err2 != nil {
			log.Println("LiangTCPClient SendMsg encrypt err: ", err2)
			return err2
		}
		buf = buf1
	}

	// write head (len of data)
	temp := make([]byte, 4)
	binary.BigEndian.PutUint32(temp, uint32(len(buf)))
	n, err := t.conn.Write(temp)
	if err != nil || n != 4 {
		log.Println("LiangTCPClient SendMsg err: ", err)
		return err
	}

	// write data
	n1, err1 := t.conn.Write(buf)
	if err1 != nil || n1 != len(buf) {
		log.Println("LiangTCPClient SendMsg err: ", err)
		return err
	}
	return nil
}

func (t *LiangTCPClient) Connect(addr string) error {

	connS, err := net.Dial("tcp", addr)
	// conn, err := net.Dial("tcp", "43.153.211.140:8081")
	if err != nil {
		log.Println("LiangTCPClient connect err: ", err)
		return err
	}
	t.conn = connS

	go t.handleConn()
	return nil
}

func (t *LiangTCPClient) handleConn() {
	defer t.conn.Close()
	log.Println("LiangTCPClient server connected: ", t.conn.RemoteAddr().String())

	// 粘包处理
	var msgBuffer bytes.Buffer

	for {
		// 单次接收
		buf := make([]byte, 2000)
		n, err := t.conn.Read(buf)
		if err != nil {
			log.Printf("LiangTCPClient read error(from %s): %s\r\n", t.conn.RemoteAddr().String(), err)
			break
		}
		// log.Printf("LiangTCPClient read %d bytes(from %s\r\n", n, conn.RemoteAddr().String())

		msgBuffer.Write(buf[:n])
		err1 := t.handleMsg(&msgBuffer)
		if err1 != nil {
			log.Printf("LiangTCP protocol error(from %s): %s\r\n", t.conn.RemoteAddr().String(), err)
			break
		}
	}
	t.OnDisconn(t.conn)
}

func (t *LiangTCPClient) handleMsg(msg *bytes.Buffer) error {
	for {
		if msg.Len() < 4 {
			// not enough for head length
			return nil
		}

		length := binary.BigEndian.Uint32(msg.Bytes())
		if int(length) > (msg.Len() - 4) {
			// not enough data
			return nil
		}

		// prepare data
		msg.Next(4)
		buf := msg.Next(int(length))

		// see AES decrypt
		if t.AESKey != nil {
			buf1, err2 := liangaes.AesCbcNoSaltDecrypt(buf, t.AESKey, t.AESKey)
			if err2 != nil {
				log.Println("LiangTCPClient RecvMsg decrypt err: ", err2)
				return err2
			}
			buf = buf1
		}

		// data is ready, call OnMsg
		err := t.OnMsg(t.conn, buf)
		if err != nil {
			return err
		}
	}
}
